// Another Simple SOCKS5/HTTP proxy server through a simple tunnel, considering a new name
// inspired by clowwindy's shadowsocks-nodejs

// let me make this clear at first: ASS is made for bypassing censorship, NOT to provide security
// ASS merely does some simple obfuscation(RC4) but NOT qualified as an encryption tool

// Links:
// node.js: http://nodejs.org/
// SOCKS5: http://tools.ietf.org/html/rfc1928

// features/known issue:
// SOCKS5 TCP CONNECT
// doesn't support SOCKS5 BIND/UDP ASSOCIATE
// HTTP proxy with CONNECT support, didn't read HTTP RFCs so compatibility might suck
// doesn't support HTTP keep-alive, so performance sucks, use SOCKS5 and/or HTTP CONNECT while you can
// IPv6 functionality not tested
// don't support any authentication method, barely any browser supports them anyway

// changes/improvements over sockstunnel:
// * backported from coffee script to js
// * updated to the non-flowing mode of the new stream interface introduced by node.js v0.10
// * full stream FSM, more robust, better compatibility
// * better RFC compliance, return some error code, not just shutdown the connection rudely
// * protocol is compatible as of now

// future plans:
// * investigate more ciphers
// * considering: apply padding on both way, to be more obfuscated, especially on the down link

// some thoughts:
// the SOCKS5 greeting is stable for a given client, e.g. 0x05(VER), 0x00(NMETHODS)
// the SOCKS5 request usually comes with the same 4 bytes: 0x05(VER), 0x01(CMD CONNECT), 0x00(RSV), 0x03(IP V4)
// together with the illy used RC4(same key for every connection), every connection is started with the same 6 bytes
// makes it very identifiable
// so ASS/sockstunnel dropped the first few bytes, shorten the SOCKS5 REQ to contain only DST.ADDR and DST.PORT
// but though, usually users initiate multiple connections to a same site(e.g. www.google.com, www.youtube.com)
// practically made more bytes of the connection predictable
// not to mention the down link usually starts with HTTP/1.1 200 OK\r\n
// a random padding could eliminate this but currently the old method is still working, so why bother

/* Version history
 * 0.1 client
 * 0.2 server
 * 0.3 client supports HTTP proxy too
 * 0.4 client supports condition
 */

var net = require("net");
var util = require("util");
var crypto = require("crypto");

// util.log("node/io.js " + process.version + " " + process.platform + " " + process.arch);

var UTF8 = "utf8";

// for the v0.10 Stream interface
var dummy = new Buffer(1);
var cipher_test = crypto.createCipher("rc4", dummy).update(dummy);
if (!(typeof cipher_test == "object") || !(cipher_test instanceof Buffer)){
	console.error("needs node v0.10.0 or higher");
	process.exit(1);
}

function  format_address(address, port){
	return address == undefined ? "UNKNOWN" : address + ":" + port;
}

function hexdump(buf, width){
	var lines = [];
	for (var i = 0; i < buf.length; i += width){
		var hexes = [];
		var chars = [];
		for (var j = i; j < i + width; ++j){
			var b = buf[j];
			if(b == undefined){
				hexes.push("  ");
			}else{
				var hex = b.toString(0x10);
				hexes.push(hex.length == 1 ? "0" + hex : hex);
				if(b >= 0x20 && b <= 0x7e){
					chars.push(String.fromCharCode(b));
				}else{
					chars.push(".");
				}
			}
		}
		lines.push(hexes.join(" ") + " | " + chars.join(""));
	}
	return lines.join("\n");
}

// TODO: improve for RFC compliance
// http://tools.ietf.org/html/rfc1035
var DOMAIN_NAME_VALIDATION_RE = /^[0-9a-zA-Z-\.]+$/;

// various states of the FSM
var STAGE_GREETING = 0;
// SOCKS5 states
var STAGE_GREETING_AUTH_METHODS = 0x101;
var STAGE_REQUEST = 0x102;
var STAGE_ALEN = 0x103;
var STAGE_DST = 0x104;
// HTTP states
var STAGE_HTTP_HEADER = 0x201;
// final states
var STAGE_REQ = 0x1000
var STAGE_PIPE = 0x1001;
var STAGE_UNPIPE = 0x2000;

var AUTH_CHOICE = new Buffer("0500", "hex");
//                                    FU CK _ CCP
var REP_GRANTED = new Buffer("05000001F55C4B20CC50", "hex");
//                                 FU CK _ FBX
var REP_0x01 = new Buffer("05010001F55C4B20FB58", "hex");
var REP_0x07 = new Buffer("05070001F55C4B20FB58", "hex");
var REP_0x08 = new Buffer("05080001F55C4B20FB58", "hex");
//                                 ADS BL O CK
var REP_0x04 = new Buffer("05040001AD");

var HTTP_200 = new Buffer("HTTP/1.1 200 OK\r\n\r\n");
var HTTP_410 = new Buffer("HTTP/1.1 204 No Content\r\nContent-Length: 0\r\n\r\n\r\n");

var CRLF = "\r\n";

var MAX_DUMMY_LEN = 0x100;
var DEFAULT_TIMEOUT = 10 * 60 * 1000;

var CODE_A = "A".charCodeAt(0);
var CODE_Z = "Z".charCodeAt(0);

function greeting_is_http(g){
	// dunno why but HTTP methods are all written in uppercase
	return g[0] >= CODE_A && g[0] <= CODE_Z && g[1] >= CODE_A && g[1] <= CODE_Z;
}

// server and client are designed in different principles:
// clients listen for LAN connections only, all clients are considered trusted
//     thus we are politely providing some error codes according to SOCKS5 RFC
//     and we do not check everything, like the current client FSM can be easily crashed with an infinite input without double CRLF, we just don't care
// servers listen for Internet connections instead
//     we drop them violently for any unexpected client behavior

function create_client(algorithm, key, host, port, cb_filter, timeout){
	if(timeout == undefined){
		timeout = DEFAULT_TIMEOUT;
	}

	var d = net.createServer(function(conn){
		var remote = null, cipher, decipher, redirect;
		var stage = STAGE_GREETING;
		// status should be carried through the entire lifespan of the connection
		// SOCKS5 status
		var nmethods;
		var atyp, alen;
		// HTTP status
		var http_buf = [];
		// status between STAGE_REQ and STAGE_PIPE
		var dst_addr, dst_port, new_req_to_dst = null, resp_to_ua = null;
		conn.on("readable", function(){
			// FSM to process SOCKS5 and/or HTTP requests
			while(true){
				var drain = false;
				var err = false;
				switch(stage){
					case STAGE_GREETING:
						var greeting_header = conn.read(2);
						if(greeting_header == null){
							drain = true;
							break;
						}
						if(greeting_header[0] == 5){ // VER
							var nmethods = greeting_header[1];
							stage = STAGE_GREETING_AUTH_METHODS;
							break;
						}else if(greeting_is_http(greeting_header)){
							http_buf.push(greeting_header[0], greeting_header[1]);
							stage = STAGE_HTTP_HEADER;
							break;
						}else{
							// invalid version
							err = true;
							return;
						}
					case STAGE_GREETING_AUTH_METHODS:
						var methods = conn.read(nmethods);
						if(methods == null){
							drain = true;
							break;
						}
						conn.write(AUTH_CHOICE);
						stage = STAGE_REQUEST;
						break;
					case STAGE_REQUEST:
						var req_header = conn.read(4);
						if(req_header == null){
							drain = true;
							break;
						}
						if(req_header[0] != 5 || req_header[2] != 0){
							// invalid version/RSV
							conn.write(REP_0x01);
							err = true;
							break;
						}
						if(req_header[1] != 1){
							// command not supported
							conn.write(REP_0x07);
							err = true;
							break;

						}
						atyp = req_header[3];
						switch(atyp){
							case 1:
							case 4:
								stage = STAGE_DST;
								break;
							case 3:
								stage = STAGE_ALEN;
								break;
							default:
								// address type not supported
								conn.write(REP_0x08);
								err = true;
								break;
						}
						break;
					case STAGE_ALEN:
						alen = conn.read(1);
						if(alen == null){
							drain = true;
							break;
						}
						alen = alen[0];
						stage = STAGE_DST;
						break;
					case STAGE_DST:
						var dst;
						switch(atyp){
							case 1:
							case 4:
								alen = atyp * 4;
								dst = conn.read(alen + 2);
								break;
							case 3:
								dst = conn.read(alen + 2);
								break;
						}
						if(dst == null){
							drain = true;
							break;
						}
						switch(atyp){
							case 1:
								dst_addr = [dst[0], dst[1], dst[2], dst[3]].join(".");
								break;
							case 3:
								dst_addr = dst.toString(UTF8, 0, alen);
								break;
							case 4:
								var h = dst.toString("hex", 0, alen);
								dst_addr = [h.slice(0, 4), h.slice(4, 8), h.slice(8, 12), h.slice(12, 16),
								     h.slice(16, 20), h.slice(20, 24), h.slice(24, 28), h.slice(28, 32)].join(":");
								break;
						}
						dst_port = dst.readUInt16BE(alen);
						util.log("ASSC " + d.address().port + ": "
							+ format_address(conn.remoteAddress, conn.remotePort)
							+ " -> " + format_address(dst_addr, dst_port));
						resp_to_ua = REP_GRANTED;
						stage = STAGE_REQ;
						break;
					case STAGE_HTTP_HEADER:
						var c = conn.read(1);
						if(c == null){
							drain = true;
							break;
						}
						http_buf.push(c[0]);
						// scan for end of header, double CRLF
						if (http_buf.length < 0x10
							|| http_buf[http_buf.length - 4] != 0x0d || http_buf[http_buf.length - 3] != 0x0a
							|| http_buf[http_buf.length - 2] != 0x0d || http_buf[http_buf.length - 1] != 0x0a){
							break;
						}
						// <CR><LF><CR><LF> found and removed
						http_buf.pop(); http_buf.pop(); http_buf.pop(); http_buf.pop();
						// we violated the RFC by parsing the header as UTF8
						var http_headers = new Buffer(http_buf).toString(UTF8).split(CRLF);
						http_buf = null;
						util.log("ASSC " + d.address().port + ": "
							+ format_address(conn.remoteAddress, conn.remotePort)
							+ " -> " + http_headers[0]);
						// console.log(hexdump(new Buffer(http_buf), 16));
						// extract dst from the request, and rewrite headers if needed
						var request_line = http_headers[0].split(" ");
						if(request_line.length < 2 || request_line.length > 3
							|| request_line[0].toUpperCase() != request_line[0]){
							err = true;
							break;
						}
						if(request_line[0] == "CONNECT"){
							// for CONNECT, we only need dst, and all headers were disregarded
							var request_target = request_line[1].split(":");
							if(request_target.length != 2){
								err = true;
								break;
							}
							dst_addr = request_target[0];
							dst_port = parseInt(request_target[1]);
							resp_to_ua = HTTP_200;
						}else{
							request_target = request_line[1];
							if(request_target.slice(0, 7).toLowerCase() != "http://"){
								err = true;
								console.error("invalid URI: " + request_target)
							}
							// I don't know how do deal with AUTH in URI, so it was dropped
							// why would anyone still using this? seriously?
							var idx_auth = request_target.indexOf("@", 7);
							var idx_path = request_target.indexOf("/", 7);
							var authority = request_target.slice(idx_auth == -1 ? 7 : idx_auth + 1,
								idx_path == -1 ? undefined : idx_path).split(":");
							dst_addr = authority[0].toLowerCase();
							dst_port = parseInt(authority[1]) || 80;
							var path = idx_path == -1 ? "*" : request_target.slice(idx_path);
							request_line[1] = path;
							// we violated the RFC by passing the version from upstream
							var new_header = [request_line.join(" ")];
							http_headers.shift();
							http_headers.forEach(function(h){
								if (h.slice(0, 6).toLowerCase() == "proxy-"){
									// discard headers begin with "proxy-";
								}else if(h.slice(0, 11).toLowerCase() == "connection:"){
									// discard the connection header, we don't support persistent connection
								}else{
									new_header.push(h);
								}
							});
							new_header.push("Connection: Close");
							new_header.push(CRLF); // for the two CRLF at end
							// console.error("new header: " + new_header.join("\n"));
							new_req_to_dst = new Buffer(new_header.join(CRLF), UTF8);
							new_header = null;
						}
						request_line = null;
						http_headers = null;
						stage = STAGE_REQ;
						break;
					case STAGE_REQ:
						if(resp_to_ua != null){
							conn.write(resp_to_ua);
						}
						resp_to_ua = null;
						redirect = cb_filter == null ? true : cb_filter(dst_addr);
						function make_pipe(remote_in, remote_out){
							// console.log(hexdump(forward_req, 16));
							// disengage the FSM and start piping
							if(new_req_to_dst != null){
								remote_in.write(new_req_to_dst);
							}
							new_req_to_dst = null;
							conn.pipe(remote_in);
							remote_out.pipe(conn);
						}
						if(redirect){
							cipher = crypto.createCipher(algorithm, key);
							decipher = crypto.createDecipher(algorithm, key);
							cipher.update(dummy);
							decipher.update(dummy);
							remote = net.connect({host: host, port: port}, function(){
								remote.pipe(decipher);
								cipher.pipe(remote);
								dst_addr = new Buffer(dst_addr, UTF8);
								var forward_req = new Buffer(1 + dst_addr.length + 2);
								forward_req.writeUInt8(dst_addr.length, 0);
								dst_addr.copy(forward_req, 1);
								forward_req.writeUInt16BE(dst_port, 1 + dst_addr.length);
								dst_addr = null;
								cipher.write(forward_req);
								forward_req = null;
								make_pipe(cipher, decipher);
							});
						}else{
							remote = net.connect({host: dst_addr, port: dst_port}, function(){
								make_pipe(remote, remote);
							});
						}
						remote.on("end", function(){
							if(stage == STAGE_PIPE){
								conn.end();
							};
						});
						remote.on("error", function(err){
							console.error("ASSC " + d.address().port + ": error occured on remote connection for "
								+ format_address(conn.remoteAddress, conn.remotePort) + ", " + err.toString());
							if(redirect){
								conn.unpipe(cipher);
								decipher.unpipe(conn);
							}
							stage = STAGE_UNPIPE;
							conn.destroy();
						});
						remote.setTimeout(timeout, function(){
							if(stage == STAGE_PIPE){
								if(redirect){
									conn.unpipe(cipher);
									decipher.unpipe(conn);
								}
								stage = STAGE_UNPIPE;
							}
							remote.end();
							conn.end();
						});
						conn.removeAllListeners("readable");
						// conn.removeAllListeners("end");
						stage = STAGE_PIPE;
						drain = true;
						break;
				} // switch(stage){ ...
				if(err){
					conn.removeAllListeners("readable");
					conn.destroy();
					break;
				}else if(drain){
					break;
				}
			} // while(true){ ...
		}); // conn.on("readable", function(){ ...

		/*
		conn.on("end", function(){
			if(stage == STAGE_PIPE){
				remote.end();
			}
		});
		*/
		conn.on("error", function(err){
			/*
			console.error("ASSC " + d.address().port + ": error occured on local connection from "
				+ format_address(conn.remoteAddress, conn.remotePort) + ", " + err.toString());
			*/
			if(stage == STAGE_PIPE){
				if(redirect){
					conn.unpipe(cipher);
					decipher.unpipe(conn);
				}
				remote.destroy();
				stage = STAGE_UNPIPE;
			}
		});
		conn.setTimeout(timeout, function(){
			if(stage == STAGE_PIPE){
				if(redirect){
					conn.unpipe(cipher);
					decipher.unpipe(conn);
				}
				remote.end();
				stage = STAGE_UNPIPE;
			}
			conn.end();
		});

	}); // net.createServer(function(conn){ ...

	d.on("error", function(err){
		console.error("ASSC " + d.address().port + ": " + err.toString());
		process.exit(2);
	});

	d.on("listening", function(){
		var addr = d.address();
		util.log("ASSC: listening on " + addr.address + ":" + addr.port + " -> " + host + ":" + port);
	})

	return d;
}

function create_server(algorithm, key, timeout){
	if(timeout == undefined){
		timeout = DEFAULT_TIMEOUT;
	}

	var d = net.createServer(function(conn){
		var cipher = crypto.createCipher(algorithm, key);
		var decipher = crypto.createDecipher(algorithm, key);
		// keep compatibility for sockstunnel
		cipher.update(dummy);
		decipher.update(dummy);

		conn.pipe(decipher);
		cipher.pipe(conn);
		var remote = null;
		var stage = STAGE_ALEN;
		var alen;
		decipher.on("readable", function(){
			// console.error("server debug: decipher readable");
			// FSM to process ASS requests
			while(true){
				var drain = false;
				var err = false;
				switch(stage){
					case STAGE_ALEN:
						alen = decipher.read(1);
						if(alen == null){
							drain = true;
							break;
						}
						alen = alen[0];
						// console.error("server debug: alen = " + alen + ", proceed to STAGE_DST");
						stage = STAGE_DST;
						break;
					case STAGE_DST:
						var dst, addr, port;
						dst = decipher.read(alen + 2);
						if(dst == null){
							drain = true;
							break;
						}
						if(dst.length != alen + 2){
							// strangely at end decipher's behavior is tricky
							// it will disregard the size parameter of .read()
							// console.error("dst.length=" + dst.length + ", alen=" + alen);
							err = true;
							break;
						}
						addr = dst.toString(UTF8, 0, alen);
						port = dst.readUInt16BE(alen);
						if(!DOMAIN_NAME_VALIDATION_RE.test(addr) || port == 0){
							console.error("server " + d.address().port
								+ ": invalid dst: alen=" + alen + ", port=" + port + ", from "
								+ format_address(conn.remoteAddress, conn.remotePort)
								+ ", dst dump:\n" + hexdump(dst, 0x10));
							err = true;
							break;
						}
						stage = STAGE_PIPE;
						util.log("ASSS " + d.address().port + ": "
							+ format_address(conn.remoteAddress, conn.remotePort)
							+ " -> " + format_address(addr, port));
						// disengage the FSM and initiate connection to dst
						drain = true;
						decipher.removeAllListeners("readable");
						// conn.removeAllListeners("end");
						remote = net.connect({host:addr, port:port}, function(){
							remote.pipe(cipher);
							decipher.pipe(remote);
						});
						remote.on("error", function(err){
							console.error("ASSS " + d.address().port
								+ ": error occured on remote connection for "
								+ format_address(conn.remoteAddress, conn.remotePort)
								+ ", " + err.toString());
							conn.unpipe(decipher);
							cipher.unpipe(conn);
							stage = STAGE_UNPIPE;
							conn.destroy();
						});
						remote.setTimeout(timeout, function(){
							if(stage == STAGE_PIPE){
								conn.unpipe(decipher);
								cipher.unpipe(conn);
								stage = STAGE_UNPIPE;
							}
							remote.end();
							conn.end();
						});
						break;

				} // switch(stage){ ...
				if(err){
					conn.removeAllListeners("readable");
					conn.destroy();
					break;
				}else if(drain){
					// console.error("server debug: drain");
					break;
				}
			} // while(true){ ...
		}); // decipher.on("readable", function(){ ...

		/*
		conn.on("end", function(){
			if(stage == STAGE_PIPE){
				remote.end();
			}
		});
		*/
		conn.on("error", function(err){
			/*
			util.log("ASSS " + d.address().port + ": error occured on local connection from "
				+ format_address(conn.remoteAddress, conn.remotePort) + ", " + err.toString());
			*/
			if(stage == STAGE_PIPE){
				conn.unpipe(decipher);
				cipher.unpipe(conn);
				remote.destroy();
				stage = STAGE_UNPIPE;
			}
		});
		conn.setTimeout(timeout, function(){
			if(stage == STAGE_PIPE){
				conn.unpipe(decipher);
				cipher.unpipe(conn);
				remote.end();
				stage = STAGE_UNPIPE;
			}
			conn.end();
		});
	}); // net.createServer(function(conn){ ...

	d.on("error", function(err){
		console.error("ASSS " + d.address().port + ": " + err.toString());
	});

	d.on("listening", function(){
		var addr = d.address();
		util.log("ASSS: listening on " + addr.address + ":" + addr.port);
	});

	return d;
} // create_server ...

function create_dump(){
	var d = net.createServer(function(conn){
		conn.on("readable", function(){
			var dump = conn.read();
			console.log(dump.length + " bytes");
			console.log(hexdump(dump, 16));
			conn.write("HTTP 200\r\nContent-Length: 3\r\n\r\nduh")
		});

		conn.on("error", function(err){
			util.log("HEXDUMP " + d.address().port + ": error occured on local connection from "
				+ format_address(conn.remoteAddress, conn.remotePort) + ", " + err.toString());
		});
	}); // net.createServer(function(conn){ ...

	d.on("listening", function(){
		var addr = d.address();
		util.log("HEXDUMP: listening on " + addr.address + ":" + addr.port);
	});

	return d;
} // create_dump ...

var DEFAULT_KEY_LEN = 48; // better use multiple of 6 for better base64 compatibility
var DEFAULT_ASS_PORT = 1984;
var DEFAULT_SOCKS5_PORT = 1080; // RFC1928
var RAN_PORT_MIN = 10000; // inclusive
var RAN_PORT_MAX = 65536; // exclusive
var LO = "127.0.0.1";

if(module.parent){
	module.exports.create_client = create_client;
}else{
	var argv = process.argv;
	if (argv.length == 9 && argv[2] == "client") {
		d = create_client(argv[3], new Buffer(argv[4], "base64"), argv[7], parseInt(argv[8]), null);
		d.listen(parseInt(argv[6]), argv[5]);
	} else if (argv.length == 7 && argv[2] == "server") {
		d = create_server(argv[3], new Buffer(argv[4], "base64"));
		d.listen(parseInt(argv[6]), argv[5]);
	}else if(argv.length == 3 && argv[2] == "test"){
		crypto.randomBytes(DEFAULT_KEY_LEN, function(ex, buf) {
			s = create_server("rc4", buf, 5 * 1000);
			s.listen(DEFAULT_ASS_PORT, LO);
			l = create_client("rc4", buf, LO, DEFAULT_ASS_PORT, null, 5 * 1000);
			l.listen(DEFAULT_SOCKS5_PORT, LO);
		});
	}else if(argv.length == 3 && argv[2] == "dump"){
		d = create_dump();
		d.listen(80, LO);
	} else if (argv.length == 3 && argv[2] == "keygen") {
		crypto.randomBytes(DEFAULT_KEY_LEN, function(ex, buf) {
			console.log(Math.floor((Math.random() * (RAND_PORT_MAX - RAND_PORT_MIN)) + RAND_PORT_MIN)
				+ " " + buf.toString("base64"));
		});
	}else if (argv.length == 4 && argv[2] == "hex2base64"){
		// sockstunnel uses hex keys, ASS uses base64
		console.log((new Buffer(argv[3], "hex")).toString("base64"));
	}else if (argv.length == 4 && argv[2] == "base642hex"){
		console.log((new Buffer(argv[3], "base64")).toString("hex"));
	}else{
		console.error("invalid arguments");
	}
}

